/*
 * KineticsSolver2.cpp
 *
 *  Created on: 28 Jul 2011
 *      Author: Allan
 */

#include "KineticSolver.h"

// GeoReact includes
#include <Utils/Algorithms.h>
#include <Utils/Assert.h>

KineticSolver::KineticSolver()
: dt_min(1.0E-6), dt_max(unsigned(-1)), algorithm(OdeAlgorithm::RungeKutta4)
{}

KineticSolver::KineticSolver(const CannonicalSystem& cannonical_system)
: dt_min(1.0E-6), dt_max(unsigned(-1)), algorithm(OdeAlgorithm::RungeKutta4)
{
	Initialiaze(cannonical_system);
}

void
KineticSolver::Initialiaze(const CannonicalSystem& cannonical_system)
{
	// Auxiliary variables
	const unsigned Nj = cannonical_system.GetNumPrimarySpecies();
	const unsigned Nk = cannonical_system.GetNumKineticSpecies();
	
	// Set the chemical system and the kinetic cannonical_system members
	this->cannonical_system = cannonical_system;
	
	// Initialiaze the ode solver
	ode_solver.Initialiaze(Nj + Nk);
	
	// Initialiaze the total equilibrium solver
	total_equilibrium_solver.Initialiaze(cannonical_system);
}

void
KineticSolver::SetMinimumTimeStep(double dt_min)
{
	this->dt_min = dt_min;
}

void
KineticSolver::SetMaximumTimeStep(double dt_max)
{
	this->dt_max = dt_max;
}

void
KineticSolver::SetAlgorithm(OdeAlgorithm algorithm)
{
	this->algorithm = algorithm;
	
	const unsigned Nj = cannonical_system.GetNumPrimarySpecies();
	const unsigned Nk = cannonical_system.GetNumKineticSpecies();
	
	ode_solver.Initialiaze(Nj + Nk, algorithm);
}

void
KineticSolver::SetControlSolution(double eps_abs, double eps_rel)
{
	ode_solver.SetControlSolution(eps_abs, eps_rel);
}

const double
KineticSolver::GetMinimumTimeStep() const
{
	return dt_min;
}

const double
KineticSolver::GetMaximumTimeStep() const
{
	return dt_max;
}

const OdeAlgorithm
KineticSolver::GetAlgorithm() const
{
	return algorithm;
}

const bool
KineticSolver::TotalEquilibrate(const VectorXd& ujnk, double T, double P, VectorXd& n)
{
	// Auxiliary variables
	const unsigned Nj = cannonical_system.GetNumPrimarySpecies();
	const unsigned Nk = cannonical_system.GetNumKineticSpecies();
	
	const VectorXd& nk = ujnk.segment(Nj, Nk);
	
	// Update the nk part of the vector (n) before the equilibrium calculation
	cannonical_system.SetCompositionKineticSpecies(nk, n);
	
	// Solve the total equilibrium problem
	ConvergenceDiagnostics diagnostics = total_equilibrium_solver.Solve(ujnk, T, P, n);
	
	Assert(diagnostics.converged, "Method did not converged.");

	return diagnostics.converged;
}

inline const double
WithinBounds(double x, double xmin, double xmax)
{
	return min(max(xmin, x), xmax);
}

const bool
KineticSolver::ReactStep(KineticStatus& status, double tend, double T, double P, VectorXd& n)
{
	// Auxiliary variables
	const unsigned Nj = cannonical_system.GetNumPrimarySpecies();
	const unsigned Nk = cannonical_system.GetNumKineticSpecies();
	
	// Set the combined vector ujnk = [uj nk]
	VectorXd ujnk(Nj + Nk);
	
	ujnk.segment(00, Nj) = cannonical_system.TotalCompositionPrimarySpecies(n);
	ujnk.segment(Nj, Nk) = cannonical_system.CompositionKineticSpecies(n);
	
	// The kinetic problem instance
	KineticProblem problem(*this, T, P, n);
	
	// Certify that the time step is within bounds (dt_min, dt_max)
	status.dt = WithinBounds(status.dt, dt_min, dt_max);
	
	// Iterate once with the kinetic problem
	bool succeeded = ode_solver.Iterate(problem, status.t, tend, status.dt, ujnk);
	
	// Certify that the time step is within bounds (dt_min, dt_max)
	status.dt = WithinBounds(status.dt, dt_min, dt_max);
	
	// Check if the current time has reached the end time
	status.finished = status.t == tend;
	
	// Check if ode iteration succeeded
	Assert(succeeded, "The ODE iteration did not succeeded.");
	
	// Improve the initial guess for the equilibrium calculation
	n = problem.n;
	
	// Solve the total equilibrium problem
	bool converged = TotalEquilibrate(ujnk, T, P, n);
	
	return converged;
}

void
KineticSolver::React(double tbegin, double tend, double T, double P, VectorXd& n)
{
	KineticStatus status;
	
	status.t  = tbegin;
	status.dt = dt_min;
	
	while(true)
	{
		ReactStep(status, tend, T, P, n);
		
		if(status.finished) break;
	}
}
